#include "MultirootNewtonRaphson.h"

MultirootNewtonRaphson::MultirootNewtonRaphson() :
solver(0), tolerance(1.0E-6), maxNumIterations(1000)
{}

MultirootNewtonRaphson::~MultirootNewtonRaphson()
{
	if(solver) gsl_multiroot_fdfsolver_free(solver);
}

void MultirootNewtonRaphson::SetNumEquations(unsigned numEquations)
{
	this->numEquations = numEquations;

	if(solver) gsl_multiroot_fdfsolver_free(solver);

	solver = gsl_multiroot_fdfsolver_alloc(gsl_multiroot_fdfsolver_gnewton, numEquations);
}

void MultirootNewtonRaphson::SetFunction(const Function& function)
{
	this->parameters.function = function;
}

void MultirootNewtonRaphson::SetJacobian(const Jacobian& jacobian)
{
	this->parameters.jacobian = jacobian;
}

void MultirootNewtonRaphson::SetFunctionJacobian(const FunctionJacobian& functionJacobian)
{
	this->parameters.functionJacobian = functionJacobian;
}

void MultirootNewtonRaphson::SetTolerance(double tolerance)
{
	this->tolerance = tolerance;
}

void MultirootNewtonRaphson::SetMaxNumIterations(unsigned maxNumIterations)
{
	this->maxNumIterations = maxNumIterations;
}

bool MultirootNewtonRaphson::Iterate(vector<double>& solution)
{
	gsl_vector* x = gsl_vector_alloc(numEquations);

	for(unsigned i = 0; i < numEquations; ++i)
		gsl_vector_set(x, i, solution[i]);

	gsl_multiroot_function_fdf functionData;

	functionData.n = numEquations;
	functionData.f = &GSL_Function;
	functionData.df = &GSL_Jacobian;
	functionData.fdf = &GSL_FunctionJacobian;
	functionData.params = &parameters;

	gsl_multiroot_fdfsolver_set(solver, &functionData, x);

	gsl_multiroot_fdfsolver_iterate(solver);

	for(unsigned i = 0; i < numEquations; ++i)
		solution[i] = gsl_vector_get(solver->x, i);

	gsl_vector_free(x);

	return (GSL_CONTINUE == gsl_multiroot_test_residual(solver->f, tolerance));
}

void MultirootNewtonRaphson::Solve(vector<double>& solution)
{
	gsl_vector* x = gsl_vector_alloc(numEquations);

	for(unsigned i = 0; i < numEquations; ++i)
		gsl_vector_set(x, i, solution[i]);

	gsl_multiroot_function_fdf functionData;

	functionData.n = numEquations;
	functionData.f = &GSL_Function;
	functionData.df = &GSL_Jacobian;
	functionData.fdf = &GSL_FunctionJacobian;
	functionData.params = &parameters;

	gsl_multiroot_fdfsolver_set(solver, &functionData, x);

	int status = GSL_CONTINUE;

	for(unsigned iter = 1; status == GSL_CONTINUE && iter <= maxNumIterations; ++iter)
	{
		status = gsl_multiroot_fdfsolver_iterate(solver);

		status = gsl_multiroot_test_residual(solver->f, tolerance);
	}

	for(unsigned i = 0; i < numEquations; ++i)
		solution[i] = gsl_vector_get(solver->x, i);

	gsl_vector_free(x);
}

int MultirootNewtonRaphson::GSL_Function(const gsl_vector* x, void* params, gsl_vector* f)
{
	Function func = ((Parameters*) params)->function;

	vector<double> xv(x->data, x->data + x->size);

	vector<double> fv(x->size);

	func(xv, fv);

	for(unsigned i = 0; i < x->size; ++i)
		gsl_vector_set(f, i, fv[i]);

	return GSL_SUCCESS;
}

int MultirootNewtonRaphson::GSL_Jacobian(const gsl_vector* x, void* params, gsl_matrix* J)
{
	Jacobian jac = ((Parameters*) params)->jacobian;

	vector<double> xv(x->data, x->data + x->size);

	vector<vector<double>> Jv(x->size, vector<double>(x->size));

	jac(xv, Jv);

	for(unsigned i = 0; i < x->size; ++i)
		for(unsigned j = 0; j < x->size; ++j)
			gsl_matrix_set(J, i, j, Jv[i][j]);

	return GSL_SUCCESS;
}

int MultirootNewtonRaphson::GSL_FunctionJacobian(const gsl_vector* x, void* params, gsl_vector* f, gsl_matrix* J)
{
	GSL_Function(x, params, f);
	GSL_Jacobian(x, params, J);

	return GSL_SUCCESS;
}
